__author__ = 'zoulida'
import  backtrader as bt
import pandas as pd
import numpy as np
import tushare as ts 
def get_data(code,start='2010-01-01',end='2020-08-31'):
    df=ts.get_k_data(code,autype='qfq',start=start,end=end)
    df.index=pd.to_datetime(df.date)
    df['openinterest']=0
    df=df[['open','high','low','close','volume','openinterest']]
    return df

#dataframe=get_data('600000',start='2015-01-01')
#dataframe.head()
#print(dataframe)

'''class TestStrategy(bt.Strategy):

    def __init__(self):
        bt.ind.MACD(self.data)
        bt.ind.MACDHisto(self.data)
        bt.ind.RSI(self.data,period=14)
        bt.ind.BBands(self.data)'''

class TradeStrategy(bt.Strategy):
    params=(('p1',12),('p2',26),('p3',9),)
    def __init__(self):
        self.order = None
        #获取MACD柱
        self.macdhist = bt.ind.MACDHisto(self.data,
                        period_me1=self.p.p1, 
                        period_me2=self.p.p2, 
                        period_signal=self.p.p3)
    def next(self):
        if not self.position:
                # 得到当前的账户价值
            total_value = self.broker.getvalue()
            #1手=100股，满仓买入
            ss=int((total_value/100)/self.datas[0].close[0])*100
            #当MACD柱大于0（红柱）且无持仓时满仓买入
            if self.macdhist > 0:
                    self.order=self.buy(size=ss)
        else:
            #当MACD柱小于0（绿柱）且持仓时全部清仓
            if self.macdhist < 0:
                    self.close()

def main(data,strategy,pf=False):
    cerebro = bt.Cerebro()
    feed = bt.feeds.PandasData(dataname=data)
    cerebro.adddata(feed) 
    #加载策略
    cerebro.addstrategy(strategy)
    # 设置初始资本为10,000
    startcash = 100000
    cerebro.broker.setcash(startcash) 
    # 设置交易手续费为 0.1%
    cerebro.broker.setcommission(commission=0.001) 
    cerebro.run()
    #获取回测结束后的总资金
    portvalue = cerebro.broker.getvalue()
    pnl = portvalue - startcash
    if pf:
        print(f'期初资金: 100000')
        print(f'总资金: {round(portvalue,2)}')
        print(f'净收益: {round(pnl,2)}')
    cerebro.plot()



data=get_data('601318','2010-03-01')
main(data,TradeStrategy,pf=True)

#期初资金：100000
#期末资金: 225440.47
#净收益: 125440.47